{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "312baf54",
   "metadata": {},
   "outputs": [],
   "source": [
    "# importing libaries\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy.io"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "aabdfc58",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the GRU model\n",
    "class GRU(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size, output_size):\n",
    "        super(GRU, self).__init__()\n",
    "\n",
    "        self.hidden_size = hidden_size\n",
    "\n",
    "        self.gru = nn.GRU(input_size, hidden_size, batch_first=True)\n",
    "        self.fc = nn.Linear(hidden_size, output_size)\n",
    "\n",
    "    def forward(self, x, hidden):\n",
    "        output, hidden = self.gru(x, hidden)\n",
    "        output = self.fc(output)\n",
    "        return output, hidden"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e0e08805",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Toy problem data\n",
    "input_size = 256  # number of columns in a dataset\n",
    "hidden_size = 32  # number of neurons\n",
    "output_size = 256  \n",
    "sequence_length = 160  # number of sequences/ number of rows\n",
    "batch_size = 1\n",
    "num_epochs = 20000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "18551192",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'__header__': b'MATLAB 5.0 MAT-file Platform: posix, Created on: Thu Jul 20 15:02:48 2023',\n",
       " '__version__': '1.0',\n",
       " '__globals__': [],\n",
       " 'h': array([[0.03949907, 0.03945555, 0.03968762, ..., 0.04081577, 0.04034067,\n",
       "         0.03995187],\n",
       "        [0.03949522, 0.03972834, 0.04007071, ..., 0.0406224 , 0.0405568 ,\n",
       "         0.04021852],\n",
       "        [0.03980837, 0.03966293, 0.04008707, ..., 0.04077249, 0.04050591,\n",
       "         0.04019281],\n",
       "        ...,\n",
       "        [0.04915863, 0.04935256, 0.04961976, ..., 0.04917885, 0.04899574,\n",
       "         0.04881741],\n",
       "        [0.04904458, 0.04916197, 0.04943338, ..., 0.04902097, 0.04897786,\n",
       "         0.04875013],\n",
       "        [0.04913692, 0.04932824, 0.04941899, ..., 0.04888785, 0.0488357 ,\n",
       "         0.04855312]])}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Load the .mat file\n",
    "v_data = scipy.io.loadmat('v.mat')\n",
    "h_data = scipy.io.loadmat('h.mat')\n",
    "x_data = scipy.io.loadmat('x.mat')\n",
    "h_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1d818fb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = x_data['X']\n",
    "u = h_data['h']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "61eddbd2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch._C.Generator at 0x7f4e0412c790>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Set random seed for reproducibility\n",
    "torch.manual_seed(45)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e95482d1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test data shape (256,)\n",
      "input data shape (160, 256)\n",
      "Target data shape (160, 256)\n"
     ]
    }
   ],
   "source": [
    "input_data = u[0:160,:]\n",
    "target_data = u[1:161, :]\n",
    "\n",
    "test_data = u[160, :]\n",
    "#test_target = u[161:201, :]\n",
    "\n",
    "print(\"test data shape\", test_data.shape)\n",
    "#print(\"test target shape\", test_target.shape)\n",
    "\n",
    "print(\"input data shape\",input_data.shape)\n",
    "print(\"Target data shape\",target_data.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "45115215",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input tensor shape torch.Size([1, 160, 256])\n",
      "Target tensor shape torch.Size([1, 160, 256])\n"
     ]
    }
   ],
   "source": [
    "# Convert data to tensors\n",
    "input_tensor = torch.tensor(input_data).view(batch_size, sequence_length, input_size).float()\n",
    "target_tensor = torch.tensor(target_data).view(batch_size, sequence_length, output_size).float()\n",
    "\n",
    "print(\"input tensor shape\",input_tensor.shape)\n",
    "print(\"Target tensor shape\",target_tensor.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "3692033e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert test data to tensors\n",
    "test_tensor = torch.tensor(test_data).view(batch_size, 1, input_size).float()\n",
    "#test_target_tensor = torch.tensor(test_target).view(batch_size, 40, output_size).float()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "592be7b8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 10/20000, Loss: 0.1309\n",
      "Epoch: 20/20000, Loss: 0.0677\n",
      "Epoch: 30/20000, Loss: 0.0440\n",
      "Epoch: 40/20000, Loss: 0.0342\n",
      "Epoch: 50/20000, Loss: 0.0320\n",
      "Epoch: 60/20000, Loss: 0.0310\n",
      "Epoch: 70/20000, Loss: 0.0306\n",
      "Epoch: 80/20000, Loss: 0.0304\n",
      "Epoch: 90/20000, Loss: 0.0304\n",
      "Epoch: 100/20000, Loss: 0.0304\n",
      "Epoch: 110/20000, Loss: 0.0304\n",
      "Epoch: 120/20000, Loss: 0.0303\n",
      "Epoch: 130/20000, Loss: 0.0303\n",
      "Epoch: 140/20000, Loss: 0.0303\n",
      "Epoch: 150/20000, Loss: 0.0303\n",
      "Epoch: 160/20000, Loss: 0.0303\n",
      "Epoch: 170/20000, Loss: 0.0303\n",
      "Epoch: 180/20000, Loss: 0.0303\n",
      "Epoch: 190/20000, Loss: 0.0303\n",
      "Epoch: 200/20000, Loss: 0.0303\n",
      "Epoch: 210/20000, Loss: 0.0303\n",
      "Epoch: 220/20000, Loss: 0.0303\n",
      "Epoch: 230/20000, Loss: 0.0303\n",
      "Epoch: 240/20000, Loss: 0.0303\n",
      "Epoch: 250/20000, Loss: 0.0303\n",
      "Epoch: 260/20000, Loss: 0.0303\n",
      "Epoch: 270/20000, Loss: 0.0304\n",
      "Epoch: 280/20000, Loss: 0.0304\n",
      "Epoch: 290/20000, Loss: 0.0303\n",
      "Epoch: 300/20000, Loss: 0.0304\n",
      "Epoch: 310/20000, Loss: 0.0304\n",
      "Epoch: 320/20000, Loss: 0.0304\n",
      "Epoch: 330/20000, Loss: 0.0303\n",
      "Epoch: 340/20000, Loss: 0.0303\n",
      "Epoch: 350/20000, Loss: 0.0303\n",
      "Epoch: 360/20000, Loss: 0.0303\n",
      "Epoch: 370/20000, Loss: 0.0303\n",
      "Epoch: 380/20000, Loss: 0.0303\n",
      "Epoch: 390/20000, Loss: 0.0304\n",
      "Epoch: 400/20000, Loss: 0.0303\n",
      "Epoch: 410/20000, Loss: 0.0303\n",
      "Epoch: 420/20000, Loss: 0.0304\n",
      "Epoch: 430/20000, Loss: 0.0304\n",
      "Epoch: 440/20000, Loss: 0.0304\n",
      "Epoch: 450/20000, Loss: 0.0304\n",
      "Epoch: 460/20000, Loss: 0.0304\n",
      "Epoch: 470/20000, Loss: 0.0304\n",
      "Epoch: 480/20000, Loss: 0.0304\n",
      "Epoch: 490/20000, Loss: 0.0304\n",
      "Epoch: 500/20000, Loss: 0.0304\n",
      "Epoch: 510/20000, Loss: 0.0304\n",
      "Epoch: 520/20000, Loss: 0.0304\n",
      "Epoch: 530/20000, Loss: 0.0304\n",
      "Epoch: 540/20000, Loss: 0.0304\n",
      "Epoch: 550/20000, Loss: 0.0305\n",
      "Epoch: 560/20000, Loss: 0.0305\n",
      "Epoch: 570/20000, Loss: 0.0305\n",
      "Epoch: 580/20000, Loss: 0.0305\n",
      "Epoch: 590/20000, Loss: 0.0304\n",
      "Epoch: 600/20000, Loss: 0.0304\n",
      "Epoch: 610/20000, Loss: 0.0304\n",
      "Epoch: 620/20000, Loss: 0.0306\n",
      "Epoch: 630/20000, Loss: 0.0304\n",
      "Epoch: 640/20000, Loss: 0.0304\n",
      "Epoch: 650/20000, Loss: 0.0304\n",
      "Epoch: 660/20000, Loss: 0.0304\n",
      "Epoch: 670/20000, Loss: 0.0304\n",
      "Epoch: 680/20000, Loss: 0.0304\n",
      "Epoch: 690/20000, Loss: 0.0304\n",
      "Epoch: 700/20000, Loss: 0.0304\n",
      "Epoch: 710/20000, Loss: 0.0304\n",
      "Epoch: 720/20000, Loss: 0.0304\n",
      "Epoch: 730/20000, Loss: 0.0305\n",
      "Epoch: 740/20000, Loss: 0.0306\n",
      "Epoch: 750/20000, Loss: 0.0305\n",
      "Epoch: 760/20000, Loss: 0.0304\n",
      "Epoch: 770/20000, Loss: 0.0304\n",
      "Epoch: 780/20000, Loss: 0.0304\n",
      "Epoch: 790/20000, Loss: 0.0305\n",
      "Epoch: 800/20000, Loss: 0.0305\n",
      "Epoch: 810/20000, Loss: 0.0305\n",
      "Epoch: 820/20000, Loss: 0.0307\n",
      "Epoch: 830/20000, Loss: 0.0305\n",
      "Epoch: 840/20000, Loss: 0.0305\n",
      "Epoch: 850/20000, Loss: 0.0306\n",
      "Epoch: 860/20000, Loss: 0.0306\n",
      "Epoch: 870/20000, Loss: 0.0307\n",
      "Epoch: 880/20000, Loss: 0.0305\n",
      "Epoch: 890/20000, Loss: 0.0305\n",
      "Epoch: 900/20000, Loss: 0.0306\n",
      "Epoch: 910/20000, Loss: 0.0304\n",
      "Epoch: 920/20000, Loss: 0.0306\n",
      "Epoch: 930/20000, Loss: 0.0305\n",
      "Epoch: 940/20000, Loss: 0.0305\n",
      "Epoch: 950/20000, Loss: 0.0305\n",
      "Epoch: 960/20000, Loss: 0.0305\n",
      "Epoch: 970/20000, Loss: 0.0305\n",
      "Epoch: 980/20000, Loss: 0.0306\n",
      "Epoch: 990/20000, Loss: 0.0306\n",
      "Epoch: 1000/20000, Loss: 0.0305\n",
      "Epoch: 1010/20000, Loss: 0.0305\n",
      "Epoch: 1020/20000, Loss: 0.0306\n",
      "Epoch: 1030/20000, Loss: 0.0307\n",
      "Epoch: 1040/20000, Loss: 0.0305\n",
      "Epoch: 1050/20000, Loss: 0.0307\n",
      "Epoch: 1060/20000, Loss: 0.0305\n",
      "Epoch: 1070/20000, Loss: 0.0305\n",
      "Epoch: 1080/20000, Loss: 0.0306\n",
      "Epoch: 1090/20000, Loss: 0.0308\n",
      "Epoch: 1100/20000, Loss: 0.0306\n",
      "Epoch: 1110/20000, Loss: 0.0307\n",
      "Epoch: 1120/20000, Loss: 0.0305\n",
      "Epoch: 1130/20000, Loss: 0.0307\n",
      "Epoch: 1140/20000, Loss: 0.0306\n",
      "Epoch: 1150/20000, Loss: 0.0305\n",
      "Epoch: 1160/20000, Loss: 0.0307\n",
      "Epoch: 1170/20000, Loss: 0.0309\n",
      "Epoch: 1180/20000, Loss: 0.0305\n",
      "Epoch: 1190/20000, Loss: 0.0306\n",
      "Epoch: 1200/20000, Loss: 0.0307\n",
      "Epoch: 1210/20000, Loss: 0.0307\n",
      "Epoch: 1220/20000, Loss: 0.0306\n",
      "Epoch: 1230/20000, Loss: 0.0305\n",
      "Epoch: 1240/20000, Loss: 0.0305\n",
      "Epoch: 1250/20000, Loss: 0.0311\n",
      "Epoch: 1260/20000, Loss: 0.0313\n",
      "Epoch: 1270/20000, Loss: 0.0306\n",
      "Epoch: 1280/20000, Loss: 0.0307\n",
      "Epoch: 1290/20000, Loss: 0.0306\n",
      "Epoch: 1300/20000, Loss: 0.0307\n",
      "Epoch: 1310/20000, Loss: 0.0304\n",
      "Epoch: 1320/20000, Loss: 0.0306\n",
      "Epoch: 1330/20000, Loss: 0.0306\n",
      "Epoch: 1340/20000, Loss: 0.0306\n",
      "Epoch: 1350/20000, Loss: 0.0308\n",
      "Epoch: 1360/20000, Loss: 0.0307\n",
      "Epoch: 1370/20000, Loss: 0.0306\n",
      "Epoch: 1380/20000, Loss: 0.0308\n",
      "Epoch: 1390/20000, Loss: 0.0306\n",
      "Epoch: 1400/20000, Loss: 0.0307\n",
      "Epoch: 1410/20000, Loss: 0.0307\n",
      "Epoch: 1420/20000, Loss: 0.0306\n",
      "Epoch: 1430/20000, Loss: 0.0308\n",
      "Epoch: 1440/20000, Loss: 0.0308\n",
      "Epoch: 1450/20000, Loss: 0.0309\n",
      "Epoch: 1460/20000, Loss: 0.0307\n",
      "Epoch: 1470/20000, Loss: 0.0307\n",
      "Epoch: 1480/20000, Loss: 0.0307\n",
      "Epoch: 1490/20000, Loss: 0.0309\n",
      "Epoch: 1500/20000, Loss: 0.0307\n",
      "Epoch: 1510/20000, Loss: 0.0307\n",
      "Epoch: 1520/20000, Loss: 0.0305\n",
      "Epoch: 1530/20000, Loss: 0.0307\n",
      "Epoch: 1540/20000, Loss: 0.0307\n",
      "Epoch: 1550/20000, Loss: 0.0307\n",
      "Epoch: 1560/20000, Loss: 0.0308\n",
      "Epoch: 1570/20000, Loss: 0.0307\n",
      "Epoch: 1580/20000, Loss: 0.0310\n",
      "Epoch: 1590/20000, Loss: 0.0308\n",
      "Epoch: 1600/20000, Loss: 0.0318\n",
      "Epoch: 1610/20000, Loss: 0.0306\n",
      "Epoch: 1620/20000, Loss: 0.0306\n",
      "Epoch: 1630/20000, Loss: 0.0306\n",
      "Epoch: 1640/20000, Loss: 0.0307\n",
      "Epoch: 1650/20000, Loss: 0.0307\n",
      "Epoch: 1660/20000, Loss: 0.0307\n",
      "Epoch: 1670/20000, Loss: 0.0307\n",
      "Epoch: 1680/20000, Loss: 0.0306\n",
      "Epoch: 1690/20000, Loss: 0.0309\n",
      "Epoch: 1700/20000, Loss: 0.0308\n",
      "Epoch: 1710/20000, Loss: 0.0308\n",
      "Epoch: 1720/20000, Loss: 0.0308\n",
      "Epoch: 1730/20000, Loss: 0.0305\n",
      "Epoch: 1740/20000, Loss: 0.0310\n",
      "Epoch: 1750/20000, Loss: 0.0308\n",
      "Epoch: 1760/20000, Loss: 0.0310\n",
      "Epoch: 1770/20000, Loss: 0.0310\n",
      "Epoch: 1780/20000, Loss: 0.0308\n",
      "Epoch: 1790/20000, Loss: 0.0309\n",
      "Epoch: 1800/20000, Loss: 0.0308\n",
      "Epoch: 1810/20000, Loss: 0.0309\n",
      "Epoch: 1820/20000, Loss: 0.0307\n",
      "Epoch: 1830/20000, Loss: 0.0319\n",
      "Epoch: 1840/20000, Loss: 0.0311\n"
     ]
    }
   ],
   "source": [
    "# Create GRU instance\n",
    "gru = GRU(input_size, hidden_size, output_size)\n",
    "\n",
    "# Loss and optimizer\n",
    "criterion = nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(gru.parameters(), lr=0.1)\n",
    "\n",
    "# Training loop\n",
    "for epoch in range(num_epochs):\n",
    "    # Set initial hidden state\n",
    "    hidden = torch.zeros(1, batch_size, hidden_size)\n",
    "\n",
    "    # Forward pass\n",
    "    output, hidden = gru(input_tensor, hidden)\n",
    "    loss = criterion(output, target_tensor)\n",
    "\n",
    "    # Backward and optimize\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "    # Print progress\n",
    "    if (epoch+1) % 10 == 0:\n",
    "        print(f'Epoch: {epoch+1}/{num_epochs}, Loss: {loss.item():.4f}')\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe1331ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(test_tensor.shape)\n",
    "prediction_tensor = torch.zeros(1, 40, 256).float()\n",
    "print(prediction_tensor.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "043049c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    hidden_pred = torch.zeros(1, batch_size, hidden_size)\n",
    "    prediction, _ = gru(test_tensor, hidden_pred)\n",
    "    prediction = prediction.view(1, 1, 256).float()\n",
    "    prediction_tensor[:, 0, :] = prediction\n",
    "    for i in range(39):\n",
    "        hidden_pred = torch.zeros(1, batch_size, hidden_size)\n",
    "        prediction, _ = gru(prediction, hidden_pred)\n",
    "        prediction = prediction.view(1, 1, 256).float()\n",
    "        prediction_tensor[:, i+1, :] = prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86c85974",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f361b56",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy.io\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "# Load the .mat file\n",
    "mat_data = scipy.io.loadmat('NLS.mat')\n",
    "\n",
    "# Following is the code to plot the data u vs x and t. u is 256*100\n",
    "# matrix. Use first 75 columns for training and 25 for testing :)\n",
    "\n",
    "# Access the variables stored in the .mat file\n",
    "# The variable names in the .mat file become keys in the loaded dictionary\n",
    "x = mat_data['x']\n",
    "t = mat_data['tt']\n",
    "u1 = mat_data['uu']\n",
    "\n",
    "# Use the loaded variables as needed\n",
    "print(x.shape)\n",
    "print(t.shape)\n",
    "print(u.shape)\n",
    "\n",
    "X, T = np.meshgrid(x, t)\n",
    "# Define custom color levels\n",
    "c_levels = np.linspace(np.min(u1), np.max(u1), 100)\n",
    "\n",
    "# Plot the contour\n",
    "plt.figure()\n",
    "plt.figure(figsize=(15, 5))\n",
    "plt.contourf(T, X, u1.T, levels=c_levels, cmap='coolwarm')\n",
    "plt.xlabel('t')\n",
    "plt.ylabel('x')\n",
    "plt.title('Schrondinger-Equation')\n",
    "plt.colorbar()  # Add a colorbar for the contour levels\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b6a7ced",
   "metadata": {},
   "outputs": [],
   "source": [
    "# true solution\n",
    "h_true = np.abs(u1)\n",
    "h_true = h_true.T\n",
    "print(h_true.shape)\n",
    "\n",
    "# exact\n",
    "u_test_full = h_true[161:201, :]\n",
    "print(u_test_full.shape)\n",
    "\n",
    "# with torch.no_grad():\n",
    "#     prediction = lem(test_tensor)\n",
    "\n",
    "# prediction.shape\n",
    "\n",
    "\n",
    "k1 = (prediction_tensor - u_test_full)**2\n",
    "u_test_full_tensor = torch.tensor(u_test_full**2)\n",
    "prediction_tensor.shape\n",
    "\n",
    "# Compute the relative L2 error norm (generalization error)\n",
    "relative_error_test = torch.mean(k1)/ torch.mean(u_test_full_tensor)\n",
    "\n",
    "print(\"Relative Error Test: \", relative_error_test.item(), \"%\")\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "125d34e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Max error\n",
    "\n",
    "R_abs = torch.max(prediction_tensor-u_test_full)\n",
    "print(R_abs)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e97cb0b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37f09dd5",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4c69b17",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdaf47fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Explained variance score\n",
    "\n",
    "import torch\n",
    "\n",
    "a = prediction_tensor\n",
    "b = u_test_full\n",
    "# Assuming 'a' is your predicted values (model's predictions) and 'b' is the true values (ground truth)\n",
    "# Make sure 'a' and 'b' are PyTorch tensors\n",
    "b = torch.tensor(b)\n",
    "# Calculate the mean of 'b'\n",
    "mean_b = torch.mean(b)\n",
    "\n",
    "# Calculate the Explained Variance Score\n",
    "numerator = torch.var(b - a)  # Variance of the differences between 'b' and 'a'\n",
    "denominator = torch.var(b)    # Variance of 'b'\n",
    "evs = 1 - numerator / denominator\n",
    "\n",
    "print(\"Explained Variance Score:\", evs.item())\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4f4b74d",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Mean absolute error\n",
    "\n",
    "R_mean = torch.mean(torch.abs(prediction_tensor - u_test_full))\n",
    "print(R_mean)\n",
    "\n",
    "## t= 1.28, 1.5, 1.37\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04ca2169",
   "metadata": {},
   "outputs": [],
   "source": [
    "### snapshots\n",
    "\n",
    "prediction_tensor = torch.squeeze(prediction_tensor)\n",
    "h = np.abs(u)\n",
    "h.shape\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b7b57a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "# Create the figure and axis objects with reduced width\n",
    "fig, ax = plt.subplots(figsize=(5, 5))  # You can adjust the width (7 inches) and height (5 inches) as needed\n",
    "\n",
    "# # Make sure the font is Times Roman\n",
    "# plt.rcParams['font.family'] = 'Times New Roman'\n",
    "\n",
    "# # Perform the prediction\n",
    "# with torch.no_grad():\n",
    "#     prediction = lem(test_tensor)\n",
    "\n",
    "final_time_output = prediction_tensor[-38, :]\n",
    "final_out = final_time_output.detach().numpy().reshape(-1, 1)\n",
    "final_true = h[-38, :].reshape(-1, 1)\n",
    "print(final_out.shape)\n",
    "print(final_true.shape)\n",
    "\n",
    "# Plot the data with red and blue lines, one with dotted and one with solid style\n",
    "ax.plot(x.T, final_out, color='red', linestyle='dotted', linewidth=12, label='Prediction')\n",
    "ax.plot(x.T, final_true, color='blue', linestyle='solid', linewidth=7, label='True')\n",
    "\n",
    "# Set the axis labels with bold font weight\n",
    "ax.set_xlabel(r\"${x}$\", fontsize=26, color='black', fontdict={'weight': 'bold'})\n",
    "ax.set_ylabel(r\"${|u(x, t)|}$\", fontsize=26, color='black', fontdict={'weight': 'bold'})\n",
    "\n",
    "# Set the title with bold font weight\n",
    "ax.set_title(r\"${t = 1.28}$\", fontsize=26, color='black', fontweight='bold')\n",
    "\n",
    "# Set the number of ticks for x-axis and y-axis to 3\n",
    "ax.set_xticks([-5, 0, 5])\n",
    "ax.set_yticks([0, 2, 4])\n",
    "\n",
    "# Set tick labels fontweight to bold and increase font size\n",
    "ax.tick_params(axis='both', which='major', labelsize=20, width=2, length=10)\n",
    "\n",
    "# # Set the fontweight for tick labels to bold\n",
    "# for tick in ax.get_xticklabels() + ax.get_yticklabels():\n",
    "#     tick.set_weight('bold')\n",
    "\n",
    "# Set the spines linewidth to bold\n",
    "ax.spines['top'].set_linewidth(2)\n",
    "ax.spines['right'].set_linewidth(2)\n",
    "ax.spines['bottom'].set_linewidth(2)\n",
    "ax.spines['left'].set_linewidth(2)\n",
    "\n",
    "\n",
    "# Increase font size for x and y axis numbers\n",
    "ax.tick_params(axis='both', which='major', labelsize=24)\n",
    "\n",
    "# Set the legend\n",
    "# ax.legend()\n",
    "\n",
    "plt.savefig('GRU_1.28_20.pdf', dpi=500, bbox_inches=\"tight\")\n",
    "\n",
    "# Show the plot\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a51774c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "# Create the figure and axis objects with reduced width\n",
    "fig, ax = plt.subplots(figsize=(5, 5))  # You can adjust the width (7 inches) and height (5 inches) as needed\n",
    "\n",
    "# # Make sure the font is Times Roman\n",
    "# plt.rcParams['font.family'] = 'Times New Roman'\n",
    "\n",
    "# # Perform the prediction\n",
    "# with torch.no_grad():\n",
    "#     prediction = lem(test_tensor)\n",
    "\n",
    "\n",
    "final_time_output = prediction_tensor[-3, :]\n",
    "final_out = final_time_output.detach().numpy().reshape(-1, 1)\n",
    "final_true = h[-3, :].reshape(-1, 1)\n",
    "print(final_out.shape)\n",
    "print(final_true.shape)\n",
    "\n",
    "# Plot the data with red and blue lines, one with dotted and one with solid style\n",
    "ax.plot(x.T, final_out, color='red', linestyle='dotted', linewidth=12, label='Prediction')\n",
    "ax.plot(x.T, final_true, color='blue', linestyle='solid', linewidth=7, label='True')\n",
    "\n",
    "# Set the axis labels with bold font weight\n",
    "ax.set_xlabel(r\"${x}$\", fontsize=26, color='black', fontdict={'weight': 'bold'})\n",
    "ax.set_ylabel(r\"${u(x, t)}$\", fontsize=26, color='black', fontdict={'weight': 'bold'})\n",
    "\n",
    "# Set the title with bold font weight\n",
    "ax.set_title(r\"${t = 1.5}$\", fontsize=26, color='black', fontweight='bold')\n",
    "\n",
    "# Set the number of ticks for x-axis and y-axis to 3\n",
    "ax.set_xticks([-5, 0, 5])\n",
    "ax.set_yticks([0, 2, 4])\n",
    "\n",
    "# Set tick labels fontweight to bold and increase font size\n",
    "ax.tick_params(axis='both', which='major', labelsize=20, width=2, length=10)\n",
    "\n",
    "# # Set the fontweight for tick labels to bold\n",
    "# for tick in ax.get_xticklabels() + ax.get_yticklabels():\n",
    "#     tick.set_weight('bold')\n",
    "\n",
    "# Set the spines linewidth to bold\n",
    "ax.spines['top'].set_linewidth(2)\n",
    "ax.spines['right'].set_linewidth(2)\n",
    "ax.spines['bottom'].set_linewidth(2)\n",
    "ax.spines['left'].set_linewidth(2)\n",
    "\n",
    "\n",
    "# Increase font size for x and y axis numbers\n",
    "ax.tick_params(axis='both', which='major', labelsize=24)\n",
    "\n",
    "# Set the legend\n",
    "# ax.legend()\n",
    "\n",
    "plt.savefig('GRU_1.5_20.pdf', dpi=500, bbox_inches=\"tight\")\n",
    "\n",
    "# Show the plot\n",
    "plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "670228bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "### 80-20 contour plot\n",
    "\n",
    "conc_u = torch.squeeze(input_tensor)\n",
    "\n",
    "concatenated_tensor = torch.cat((conc_u, prediction_tensor), dim=0)\n",
    "\n",
    "t1 = np.linspace(0, 1.5707 , 200)\n",
    "\n",
    "import torch\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.ticker import FixedLocator\n",
    "\n",
    "# Assuming you have defined concatenated_tensor as a PyTorch tensor\n",
    "# concatenated_tensor = torch.cat((tensor1, tensor2), dim=0)\n",
    "\n",
    "# Convert concatenated_tensor to a NumPy array\n",
    "concatenated_array = concatenated_tensor.numpy()\n",
    "\n",
    "# Define custom color levels\n",
    "x = np.linspace(-5, 5, concatenated_array.shape[1])  # Replace 0 and 1 with your actual x range\n",
    "t = np.linspace(0, 1.57, concatenated_array.shape[0])  # Replace 0 and 1 with your actual t range\n",
    "X, T = np.meshgrid(x, t1)\n",
    "\n",
    "# Define custom color levels using the minimum and maximum from the NumPy array\n",
    "c_levels = np.linspace(np.min(concatenated_array), np.max(concatenated_array), 400)\n",
    "\n",
    "# Plot the contour with interpolated data\n",
    "plt.figure(figsize=(20, 5))\n",
    "plt.pcolormesh(T, X, concatenated_array, shading='auto', cmap='twilight')\n",
    "\n",
    "# Set the fontweight for axis labels to regular (not bold)\n",
    "plt.xlabel(\"$t$\", fontsize=26)\n",
    "plt.ylabel(\"$x$\", fontsize=26)\n",
    "plt.title(\"$|u(x, t)|$\", fontsize=26)\n",
    "\n",
    "# Set tick labels fontweight to regular (not bold) and increase font size\n",
    "plt.tick_params(axis='both', which='major', labelsize=20, width=3, length=10)\n",
    "\n",
    "# Set the fontweight for tick labels to regular (not bold)\n",
    "for tick in plt.gca().get_xticklabels() + plt.gca().get_yticklabels():\n",
    "    tick.set_weight('normal')\n",
    "\n",
    "# Set the number of ticks for x-axis and y-axis to 5\n",
    "num_ticks = 5\n",
    "x_ticks = np.linspace(np.min(T), np.max(T), num_ticks)\n",
    "y_ticks = np.linspace(np.min(X), np.max(X), num_ticks)\n",
    "\n",
    "plt.gca().xaxis.set_major_locator(FixedLocator(x_ticks))\n",
    "plt.gca().yaxis.set_major_locator(FixedLocator(y_ticks))\n",
    "\n",
    "cbar1 = plt.colorbar()\n",
    "# Set the number of ticks for the color bar with uniformly distributed numbers\n",
    "num_ticks = 5\n",
    "c_ticks = np.linspace(np.min(concatenated_array), np.max(concatenated_array), num_ticks)\n",
    "cbar1.set_ticks(c_ticks)\n",
    "\n",
    "# Set the fontweight and fontsize for color bar tick labels\n",
    "for t in cbar1.ax.get_yticklabels():\n",
    "    t.set_weight('normal')\n",
    "    t.set_fontsize(26)  # Increase the font size for color bar tick labels\n",
    "\n",
    "# Increase the size of numbers on axis and color bar\n",
    "plt.xticks(fontsize=26)\n",
    "plt.yticks(fontsize=26)\n",
    "\n",
    "# Increase the tick size and width of the color bar\n",
    "cbar1.ax.tick_params(axis='both', which='major', labelsize=30, width=3,  length=10)\n",
    "\n",
    "# Add a dotted line at t = 0.8\n",
    "plt.axvline(x=1.26449, color='black', linestyle='dotted', linewidth=5)\n",
    "\n",
    "#plt.savefig('Contour_LEM_20.pdf', dpi=500, bbox_inches=\"tight\")\n",
    "plt.savefig('contour_GRU_20.jpeg', dpi=500, bbox_inches=\"tight\")\n",
    "# Show the plot\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8b48c55",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bff925ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "# Create the figure and axis objects with reduced width\n",
    "fig, ax = plt.subplots(figsize=(5, 5))  # You can adjust the width (7 inches) and height (5 inches) as needed\n",
    "\n",
    "# # Make sure the font is Times Roman\n",
    "# plt.rcParams['font.family'] = 'Times New Roman'\n",
    "\n",
    "# # Perform the prediction\n",
    "# with torch.no_grad():\n",
    "#     prediction = lem(test_tensor)\n",
    "\n",
    "\n",
    "final_time_output = prediction_tensor[-23, :]\n",
    "final_out = final_time_output.detach().numpy().reshape(-1, 1)\n",
    "final_true = h[-23, :].reshape(-1, 1)\n",
    "print(final_out.shape)\n",
    "print(final_true.shape)\n",
    "\n",
    "# Plot the data with red and blue lines, one with dotted and one with solid style\n",
    "ax.plot(x.T, final_out, color='red', linestyle='dotted', linewidth=12, label='Prediction')\n",
    "ax.plot(x.T, final_true, color='blue', linestyle='solid', linewidth=7, label='True')\n",
    "\n",
    "# Set the axis labels with bold font weight\n",
    "ax.set_xlabel(r\"${x}$\", fontsize=26, color='black', fontdict={'weight': 'bold'})\n",
    "ax.set_ylabel(r\"${u(x, t)}$\", fontsize=26, color='black', fontdict={'weight': 'bold'})\n",
    "\n",
    "# Set the title with bold font weight\n",
    "ax.set_title(r\"${t = 1.5}$\", fontsize=26, color='black', fontweight='bold')\n",
    "\n",
    "# Set the number of ticks for x-axis and y-axis to 3\n",
    "ax.set_xticks([-5, 0, 5])\n",
    "ax.set_yticks([0, 2, 4])\n",
    "\n",
    "# Set tick labels fontweight to bold and increase font size\n",
    "ax.tick_params(axis='both', which='major', labelsize=20, width=2, length=10)\n",
    "\n",
    "# # Set the fontweight for tick labels to bold\n",
    "# for tick in ax.get_xticklabels() + ax.get_yticklabels():\n",
    "#     tick.set_weight('bold')\n",
    "\n",
    "# Set the spines linewidth to bold\n",
    "ax.spines['top'].set_linewidth(2)\n",
    "ax.spines['right'].set_linewidth(2)\n",
    "ax.spines['bottom'].set_linewidth(2)\n",
    "ax.spines['left'].set_linewidth(2)\n",
    "\n",
    "\n",
    "# Increase font size for x and y axis numbers\n",
    "ax.tick_params(axis='both', which='major', labelsize=24)\n",
    "\n",
    "# Set the legend\n",
    "# ax.legend()\n",
    "\n",
    "\n",
    "# Show the plot\n",
    "plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ba80a0f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12d7d2b2",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "language": "python",
   "name": "pytorch"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
